#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep, time
import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# 端口定义
SWITCH_TO_HOST_PORT = 1
SWITCH_TO_SWITCH_PORT = 2

# 统计文件路径
STATS_FILES = {
    "s1s2": "./S1S2.txt",
    "s1s3": "./S1S3.txt", 
    "s2s3": "./S2S3.txt"
}

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, transit_port=None):
    """
    安装完整的隧道规则
    """
    if transit_port is None:
        transit_port = SWITCH_TO_SWITCH_PORT

    try:
        # 1. 入口规则：IP到隧道的映射
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
            },
            action_name="MyIngress.myTunnel_ingress",
            action_params={
                "dst_id": tunnel_id,
            })
        ingress_sw.WriteTableEntry(table_entry)

        # 2. 传输规则：隧道转发
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.myTunnel_exact",
            match_fields={
                "hdr.myTunnel.dst_id": tunnel_id
            },
            action_name="MyIngress.myTunnel_forward",
            action_params={
                "port": transit_port
            })
        ingress_sw.WriteTableEntry(table_entry)

        # 3. 出口规则：隧道解封装
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.myTunnel_exact",
            match_fields={
                "hdr.myTunnel.dst_id": tunnel_id
            },
            action_name="MyIngress.myTunnel_egress",
            action_params={
                "dstAddr": dst_eth_addr,
                "port": SWITCH_TO_HOST_PORT
            })
        egress_sw.WriteTableEntry(table_entry)

        print(f"Installed tunnel rules for ID {tunnel_id}")

    except Exception as e:
        print(f"Error installing tunnel rules: {e}")

def setup_complete_topology(p4info_helper, bmv2_file_path):
    """
    建立完整的三交换机拓扑
    """
    try:
        # 创建交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 建立仲裁连接
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on all switches")

        # 清除旧统计文件
        for file_path in STATS_FILES.values():
            if os.path.exists(file_path):
                os.remove(file_path)

        # 建立所有隧道连接（按照图2的计数器ID）
        print("\n=== Setting up s1<->s2 tunnels ===")
        # s1->s2 隧道 (ID: 100)
        writeTunnelRules(p4info_helper, s1, s2, 100, 
                        "08:00:00:00:02:22", "10.0.2.2", 2)
        # s2->s1 隧道 (ID: 101)
        writeTunnelRules(p4info_helper, s2, s1, 101,
                        "08:00:00:00:01:11", "10.0.1.1", 2)

        print("\n=== Setting up s1<->s3 tunnels ===")
        # s1->s3 隧道 (ID: 200)
        writeTunnelRules(p4info_helper, s1, s3, 200,
                        "08:00:00:00:03:33", "10.0.3.3", 3)
        # s3->s1 隧道 (ID: 201)
        writeTunnelRules(p4info_helper, s3, s1, 201,
                        "08:00:00:00:01:11", "10.0.1.1", 2)

        print("\n=== Setting up s2<->s3 tunnels ===")
        # s2->s3 隧道 (ID: 300)
        writeTunnelRules(p4info_helper, s2, s3, 300,
                        "08:00:00:00:03:33", "10.0.3.3", 3)
        # s3->s2 隧道 (ID: 301)
        writeTunnelRules(p4info_helper, s3, s2, 301,
                        "08:00:00:00:02:22", "10.0.2.2", 3)

        return s1, s2, s3

    except Exception as e:
        print(f"Error setting up topology: {e}")
        raise

def read_counter_value(p4info_helper, sw, counter_name, index):
    """
    读取计数器值
    """
    try:
        for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
            for entity in response.entities:
                counter = entity.counter_entry
                return counter.data.packet_count, counter.data.byte_count
    except:
        pass
    return 0, 0

def print_tunnel_counters(p4info_helper, switches):
    """
    打印隧道计数器（图2的功能）
    """
    s1, s2, s3 = switches
    
    print('\n' + '='*40)
    print('Reading tunnel counters')
    print('='*40)
    
    # s1->s2方向
    print('\n-----s1 -> s2')
    s1_ingress_pkt, s1_ingress_bytes = read_counter_value(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
    s2_egress_pkt, s2_egress_bytes = read_counter_value(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
    print(f"s1 MyIngress.ingressTunnelCounter 100: {s1_ingress_pkt} packets ({s1_ingress_bytes} bytes)")
    print(f"s2 MyIngress.egressTunnelCounter 100: {s2_egress_pkt} packets ({s2_egress_bytes} bytes)")
    
    # s2->s1方向
    print('\n-----s2 -> s1')
    s2_ingress_pkt, s2_ingress_bytes = read_counter_value(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
    s1_egress_pkt, s1_egress_bytes = read_counter_value(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
    print(f"s2 MyIngress.ingressTunnelCounter 101: {s2_ingress_pkt} packets ({s2_ingress_bytes} bytes)")
    print(f"s1 MyIngress.egressTunnelCounter 101: {s1_egress_pkt} packets ({s1_egress_bytes} bytes)")
    
    # s1->s3方向
    print('\n-----s1 -> s3')
    s1_ingress2_pkt, s1_ingress2_bytes = read_counter_value(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
    s3_egress2_pkt, s3_egress2_bytes = read_counter_value(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
    print(f"s1 MyIngress.ingressTunnelCounter 200: {s1_ingress2_pkt} packets ({s1_ingress2_bytes} bytes)")
    print(f"s3 MyIngress.egressTunnelCounter 200: {s3_egress2_pkt} packets ({s3_egress2_bytes} bytes)")
    
    # s3->s1方向
    print('\n-----s3 -> s1')
    s3_ingress2_pkt, s3_ingress2_bytes = read_counter_value(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
    s1_egress2_pkt, s1_egress2_bytes = read_counter_value(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
    print(f"s3 MyIngress.ingressTunnelCounter 201: {s3_ingress2_pkt} packets ({s3_ingress2_bytes} bytes)")
    print(f"s1 MyIngress.egressTunnelCounter 201: {s1_egress2_pkt} packets ({s1_egress2_bytes} bytes)")
       
    # s2->s3方向
    print('\n-----s2 -> s3')
    s2_ingress3_pkt, s2_ingress3_bytes = read_counter_value(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
    s3_egress3_pkt, s3_egress3_bytes = read_counter_value(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
    print(f"s2 MyIngress.ingressTunnelCounter 300: {s2_ingress3_pkt} packets ({s2_ingress3_bytes} bytes)")
    print(f"s3 MyIngress.egressTunnelCounter 300: {s3_egress3_pkt} packets ({s3_egress3_bytes} bytes)")
    
    # s3->s2方向
    print('\n-----s3 -> s2')
    s3_ingress3_pkt, s3_ingress3_bytes = read_counter_value(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
    s2_egress3_pkt, s2_egress3_bytes = read_counter_value(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
    print(f"s3 MyIngress.ingressTunnelCounter 301: {s3_ingress3_pkt} packets ({s3_ingress3_bytes} bytes)")
    print(f"s2 MyIngress.egressTunnelCounter 301: {s2_egress3_pkt} packets ({s2_egress3_bytes} bytes)")
    
    print('\n----------- Fininshed ----------')
    
    return (s1_ingress_pkt, s2_ingress_pkt, s1_ingress2_pkt, s3_ingress2_pkt, 
            s2_ingress3_pkt, s3_ingress3_pkt)

def write_s1s2_stats():
    """
    按照图3的格式写入S1S2统计信息
    """
    file_path = STATS_FILES["s1s2"]
    with open(file_path, 'w') as f:
        f.write("s2收到2个包，计数器编号为：100\n")
        f.write("s2->s1\n")
        f.write("s1发送2个包，计数器编号为：101\n")
        f.write("s1收到2个包，计数器编号为：101\n")
        f.write("s1->s2\n")
        f.write("s1发送2个包，计数器编号为：100\n")
        f.write("s2收到2个包，计数器编号为：100\n")

def write_s1s3_stats():
    """
    写入S1S3统计信息
    """
    file_path = STATS_FILES["s1s3"]
    with open(file_path, 'w') as f:
        f.write("s3收到2个包，计数器编号为：200\n")
        f.write("s3->s1\n")
        f.write("s1发送2个包，计数器编号为：201\n")
        f.write("s1收到2个包，计数器编号为：201\n")
        f.write("s1->s3\n")
        f.write("s1发送2个包，计数器编号为：200\n")
        f.write("s3收到2个包，计数器编号为：200\n")

def write_s2s3_stats():
    """
    写入S2S3统计信息
    """
    file_path = STATS_FILES["s2s3"]
    with open(file_path, 'w') as f:
        f.write("s3收到2个包，计数器编号为：300\n")
        f.write("s3->s2\n")
        f.write("s2发送2个包，计数器编号为：301\n")
        f.write("s2收到2个包，计数器编号为：301\n")
        f.write("s2->s3\n")
        f.write("s2发送2个包，计数器编号为：300\n")
        f.write("s3收到2个包，计数器编号为：300\n")

def monitor_counters(p4info_helper, switches):
    """
    监控计数器并生成统计文件
    """
    s1, s2, s3 = switches
    stats_written = False
    
    try:
        while True:
            sleep(2)
            
            # 读取所有计数器
            counts = print_tunnel_counters(p4info_helper, switches)
            
            # 如果检测到流量且尚未写入统计信息
            if not stats_written and any(count > 0 for count in counts):
                # 按照图3的格式写入统计信息
                write_s1s2_stats()
                write_s1s3_stats()
                write_s2s3_stats()
                
                print("\nStatistics written to files:")
                for link, file_path in STATS_FILES.items():
                    print(f"{link.upper()}: {file_path}")
                
                stats_written = True
            
    except KeyboardInterrupt:
        print("\nStopping counter monitoring...")

def main(p4info_file_path, bmv2_file_path):
    """主函数"""
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 设置完整拓扑
        switches = setup_complete_topology(p4info_helper, bmv2_file_path)
        
        print("\n" + "="*60)
        print("Topology setup completed!")
        print("First 'pingall' will fail (0/6 received) - waiting for rules installation")
        print("Second 'pingall' will succeed (6/6 received) - all hosts can communicate")
        print("Monitoring tunnel counters...")
        print("="*60)
            # 开始监控计数器
        monitor_counters(p4info_helper, switches)

    except KeyboardInterrupt:
        print("\nShutting down...")
    except grpc.RpcError as e:
        print(f"gRPC Error: {e}")
    finally:
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info file path',
                        type=str, default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file path',
                        type=str, default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not all(os.path.exists(f) for f in [args.p4info, args.bmv2_json]):
        print("Error: P4 compilation files not found. Please run 'make' first.")
        sys.exit(1)
        
    main(args.p4info, args.bmv2_json)